package cryption

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"math/big"
)

// 初始化向量
var aesIV []byte = []byte("1234567887654321")
var AesKey = []byte("12a45b78e765w3j1")

// 使用PKCS7进行填充，IOS也是7
func PKCS7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func PKCS7UnPadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

func PKCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func PKCS5UnPadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

func AesCBCEncrypt(rawData, key []byte) ([]byte, error) {
	// aes加密, 填充秘钥key的16位, 24,32分别对应AES-128, AES-192, or AES-256.
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 填充原文
	blockSize := block.BlockSize()
	rawData = PKCS7Padding(rawData, blockSize)

	cipherText := make([]byte, len(rawData))
	mode := cipher.NewCBCEncrypter(block, aesIV)
	mode.CryptBlocks(cipherText, rawData)

	return cipherText, nil
}

func AesCBCDecrypt(encryptData, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	blockSize := block.BlockSize()

	if len(encryptData) < blockSize {
		return nil, errors.New("ciphertext too short")
	}

	// CBC 模式始终适用于整个块。
	if len(encryptData)%blockSize != 0 {
		return nil, errors.New("ciphertext is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, aesIV)

	origData := make([]byte, len(encryptData))
	mode.CryptBlocks(origData, encryptData)

	// 解填充
	origData = PKCS7UnPadding(origData)
	return origData, nil
}

/**
 * @description: AES加密
 * @param {string} rawData 需要加密的数据
 * @param {string} key AES密钥
 * @return {string, error} 加密后数据
 */
func AesEncrypt(rawData, key string) (string, error) {
	data, err := AesCBCEncrypt([]byte(rawData), []byte(key))
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

/**
 * @description: AES解密
 * @param {string} rawData 需要解密的数据
 * @param {string} key AES密钥
 * @return {string, error} 解密后数据
 */
func AesDecrypt(rawData, key []byte) (string, error) {
	data, err := base64.StdEncoding.DecodeString(string(rawData))
	if err != nil {
		return "", err
	}
	dnData, err := AesCBCDecrypt(data, key)
	if err != nil {
		return "", err
	}
	return string(dnData), nil
}

/**
 * @description: 生成随机字符串
 * @param {int} n 长度
 * @param {...[]rune} allowedChars 允许的字符
 * @return {string}
 */
func RandomString(n int, allowedChars ...[]rune) string {
	defaultLetters := []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")

	var letters []rune

	if len(allowedChars) == 0 {
		letters = defaultLetters
	} else {
		letters = allowedChars[0]
	}

	b := make([]rune, n)
	for i := range b {
		len := int64(len(letters))

		x, _ := rand.Int(rand.Reader, big.NewInt(len))

		b[i] = letters[x.Int64()]
	}

	return string(b)
}
